Skip to content

Add defensive load_weights helper to normalise calibration weight shape#351

Merged
MaxGhenis merged 1 commit into
mainfrom
fix/weight-shape-load-helper
Apr 17, 2026
Merged

Add defensive load_weights helper to normalise calibration weight shape#351
MaxGhenis merged 1 commit into
mainfrom
fix/weight-shape-load-helper

Conversation

@MaxGhenis

Copy link
Copy Markdown
Contributor

Summary

The bug hunt (finding U4) flagged that the L2 calibrator in calibrate_local_areas writes 2D (n_areas, n_records) weights under a dataset key, while the L0-regularised variant (in a separate PR branch) writes a flat (n_records,) array under the same key. Downstream consumers cannot tell the shapes apart unless they explicitly check axes; .sum(axis=0) on a 1D array is a no-op, which silently produces wrong answers.

This PR adds a shape-normalising loader in the stable module so any future consumer — including whichever calibrator ships first — can load weights safely:

  • load_weights(weight_file, dataset_key, n_areas=None, n_records=None) returns a 2D (n_areas, n_records) array. A 1D input is reshaped to (1, n_records) so matrix ops behave identically across back-ends.
  • Optional n_areas and n_records kwargs raise a clear ValueError on shape mismatch instead of letting the error surface as garbage weighted totals downstream.
  • Missing dataset keys raise KeyError with the list of available keys. 3D+ inputs raise ValueError explicitly.

No call site is changed here — this is a detection and normalisation tool that consumers can adopt as they are touched (e.g. matrix_builder, validation notebooks). The L0 calibrator PR should route through this helper before merging to main so both shapes stay interchangeable.

Test plan

  • uv run pytest policyengine_uk_data/tests/test_load_weights.py -q passes (5 cases: 2D pass-through, 1D promotion, explicit shape validation, missing key, higher-dim rejection).

Finding U4 from the bug hunt.

@MaxGhenis

Copy link
Copy Markdown
Contributor Author

Self-review: APPROVE. See /tmp/bug-hunt/reviews/uk-data-reviews.md for detail. (GitHub blocks self-approve, so leaving as comment.)

@MaxGhenis MaxGhenis merged commit a87047d into main Apr 17, 2026
3 checks passed
@MaxGhenis MaxGhenis deleted the fix/weight-shape-load-helper branch April 17, 2026 16:05
vahid-ahmadi added a commit that referenced this pull request Apr 21, 2026
Conflict was in `policyengine_uk_data/utils/calibrate.py`: main added
`load_weights` (#351 defensive h5-weight loader) at the same file
position where this branch added `compute_log_weight_smoothness_penalty`
(#346 step 5). Both functions are independent and both stay.

All 150 tests in the panel-pipeline suite + the calibrate smoothness
tests pass after the merge, including `load_weights` consumers pulled
in from main (test_calibrate_save, test_la_land_value_targets).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant